#!/usr/bin/env python
# coding=utf-8

"""
@author: xu
@license: MIT
@file: gribfile.py
@date: 2023/11/1 21:26
@desc: 
"""
from datetime import datetime, timedelta

import pygrib
import os
import sys
import numpy as np
from numpy.ma import MaskedArray
from loguru import logger
from unitxu.draw.grid2png import array2draw, array2tiff, array2png


class Grib:
    """
    单通道处理类
    """
    def __init__(self, grib, beishu=1, latRange=None, lonRange=None, fillvalue=99999, args=[], **kwargs):
        self.grib = grib
        self.beishu = beishu
        self.fillvalue = fillvalue
        self.latRange = latRange
        self.lonRange = lonRange
        self.get_data()
        self.args = args
        self.kwargs = kwargs
        self.get_datameta()
        self.data_slice_dicts = {}
        # del self.grib

    def upper_data(self):
        maxdata = np.nanmax(self.datas)
        if isinstance(self.datas, MaskedArray):
            self.fillvalue = np.array(self.datas[self.datas.mask])[0]
            self.datas = np.array(self.datas)
        maxdata = maxdata if maxdata > 0 else np.nanmax(-self.datas)
        for i in range(3, -6, -1):
            if 10 ** (i - 1) < maxdata <= 10 ** i:
                self.beishu = 10 ** (-i + 4)
                break
            else:
                self.beishu = 1
        self.datas = np.round(self.datas * self.beishu, 0)

    def get_data(self):
        if self.latRange is None and self.lonRange is None:
            self.lats, self.lons = self.grib.latlons()
            self.datas = self.grib.values
        else:
            self.datas, self.lats, self.lons = self.grib.data(lat1=self.latRange[0], lat2=self.latRange[1],
                                                              lon1=self.lonRange[0], lon2=self.lonRange[1])
        if self.beishu != 1:
            self.upper_data()
        if isinstance(self.datas, MaskedArray):
            self.fillvalue = np.array(self.datas[self.datas.mask])[0]
            self.datas = np.array(self.datas)


    def get_datameta(self):
        """
        必须要在get_data方法之后调用
        :return:
        """
        rownums, locnums = self.lats.shape
        self.latmin, self.latmax = np.nanmin(self.lats), np.nanmax(self.lats)
        self.lonmin, self.lonmax = np.nanmin(self.lons), np.nanmax(self.lons)
        self.latapi = (self.latmax - self.latmin) / (rownums - 1)
        self.lonapi = (self.lonmax - self.lonmin) / (locnums - 1)
        self.name = self.grib.name
        self.shortname = self.kwargs.get(self.grib.name, self.grib.shortName.upper())
        try:
            if 'P1' in self.grib.keys():
                fcst_time = self.grib.P1
            else:
                fcst_time = self.grib.__str__().split(':')[-2].strip().split()[2]
                fcst_time = fcst_time.split('-')[-1]
                fcst_time = int(fcst_time)
        except:
            fcst_time = 0
        self.qibaotime = '%s-%02d-%02d %02d:00:00' % (self.grib.year, self.grib.month, self.grib.day, self.grib.hour)
        self.yubaotime = datetime.strptime(self.qibaotime, '%Y-%m-%d %H:00:00') + timedelta(hours=fcst_time)
        self.yubaotime = self.yubaotime.strftime('%Y-%m-%d %H:00:00')
        self.unit = self.grib.parameterUnits
        for a in self.args:
            try:
                self.__setattr__(a, self.grib.__getattribute__(a))
            except:
                print(a, 'waring不存在')
        # self.level = self.grib.level
        self.message = self.grib.__str__()
        self.level = self.message.split(':')[5].replace('levels', '').replace('level', '').strip()
        self.metadata = {
            'target': self.shortname, 'unit': self.unit, 'level': self.level, 'qibaotime': self.qibaotime,
            'yubaotime': self.yubaotime
        }

    def data_slice(self, lat_split_nums=None, lon_split_nums=None):
        """
        :param lat_split_nums: 一个单元块纵向包含的子单元个数
        :param lon_split_nums: 一个单元块横向包含的子单元个数
        :return: {'block_0_0': 2darray, 'block_0_10': 2darray....}
        """
        # import sharppy.sharptab.profile as profile
        # import sharppy.sharptab.params as params
        # prof = profile.create_profile(date=cur_date, profile='convective', pres=press, hght=hght, tmpc=tmpc, \
        #                               dwpc=dwpc, wdir=wind, wspd=wins, missing=-9999.0, strictQC=False)
        # ki = params.k_index(prof)
        if lat_split_nums is None:
            self.data_slice_dicts = {'block_0_0': self.datas.tolist()}
            self.lat_split_nums, self.lon_split_nums = self.datas.shape
        else:
            self.lat_split_nums, self.lon_split_nums = lat_split_nums, lon_split_nums
            rownums, locnums = self.datas.shape
            # 横纵被切分了多少个单元快
            lat_slice, lon_slice = int(np.ceil(rownums / lat_split_nums)), int(np.ceil(locnums / lon_split_nums))
            for lat_index in range(lat_slice):
                for lon_index in range(lon_slice):
                    self.data_slice_dicts['block_%s_%s' % (lat_index * lat_split_nums, lon_index * lon_split_nums)] \
                        = self.datas[lat_split_nums * lat_index: lat_split_nums*lat_index + lat_split_nums,
                                     lon_index * lon_split_nums: lon_index*lon_split_nums + lon_split_nums].tolist()



class Gribfile:
    def __init__(self, file, beishu=1, fillvalue=99999, levels=None, latlons=None, output_folder=None, **kwargs):
        self.set_target(**kwargs)
        self.levels = levels
        self.latlons = latlons
        self.file = file
        self.grs = pygrib.open(self.file)
        self.kwargs = kwargs
        self.beishu = beishu
        self.fillvalue = fillvalue
        self.output_folder = output_folder

    def show_grib(self):
        for i in self.grs:
            logger.info(i)


    def get_data(self, levels=None, targets=None, typeOfLevel=None):
        """通过高度、要素、高度类型返回指定数据数组"""
        lists = []
        if typeOfLevel is None:
            grbs = self.grs.select(name=targets, level=levels)
        else:
            grbs = self.grs.select(name=targets, level=levels, typeOfLevel=typeOfLevel)
        for grb in grbs:
            if isinstance(self.latlons, list) and len(self.latlons) == 4:
                g = Grib(grb, latRange=self.latlons[:2], lonRange=self.latlons[2:],
                     beishu=self.beishu, fillvalue=self.fillvalue, kwargs=self.kwargs)
            else:
                g = Grib(grb, beishu=self.beishu, fillvalue=self.fillvalue, kwargs=self.kwargs)
            lists.append(g.datas)
        return np.array(lists)


    def get_gribs(self):
        """获取所有要素或target_dict的指定的要素数据集，列表元素为Grid对象"""
        lists = []
        if self.target_dict:
            grs = self.grs.select(name=list(self.target_dict.keys()))
        else:
            grs = self.grs
        for grb in grs:
            if isinstance(self.latlons, list) and len(self.latlons) == 4:
                lists.append(Grib(grb, latRange=self.latlons[:2], lonRange=self.latlons[2:],
                                  beishu=self.beishu, fillvalue=self.fillvalue, kwargs=self.kwargs))
            else:
                lists.append(Grib(grb, beishu=self.beishu, fillvalue=self.fillvalue, kwargs=self.kwargs))
        return lists


    def get_savename(self, target, dims, savetype=1):
        dims_format = []
        for d in dims:
            if isinstance(d, str):
                temp = d.replace(' ', '').replace('-', '').replace(':', '')
                dims_format.append({temp: temp}.get(temp, '999'))
            else:
                dims_format.append(str(round(d, 4)).replace('.', 'p'))
        output = os.path.join(self.output_folder, os.sep.join(dims_format))
        if not os.path.exists(output):
            os.makedirs(output)
        if savetype == 1:
            tiff = os.path.join(output, target + '.png')
        else:
            tiff = os.path.join(output, target + '.Tiff')
        return tiff


    def resolve_target(self, savetype=1, output_folder=None):
        """savetype 1保存为图片 2为tiff"""
        grbs = self.get_gribs()
        if output_folder is not None:
            self.output_folder = output_folder
        if savetype == 2:
            self.draw = array2tiff
        elif savetype == 1:
            self.draw = array2draw
        for grb in grbs:
            savename = self.get_savename(grb.shortname, [grb.qibaotime, grb.yubaotime,
                                          grb.level], savetype)
            self.draw([grb.latmin, grb.latmax, grb.lonmin, grb.lonmax], grb.datas, savename,
                      grb.metadata,
                      pro_map=None)


    def set_target(self, **kwargs):
        self.target_dict = kwargs